import paramiko
from io import  StringIO
import threading
import queue
class SSH():
    def __init__(self, ip, port, username, password=None, key=None, timeout=6):
        self.ip = ip
        self.port = port
        self.username = username
        self.password = password
        self.key = key
        self.timeout = timeout

    def command(self, command):
        # 绑定实例
        ssh = paramiko.SSHClient()
        # AutoAddPolicy()自动添加主机keys
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            if self.password:
                ssh.connect(hostname=self.ip, port=self.port, username=self.username, password=self.password, timeout=self.timeout)
            else:
                cache = StringIO(self.key) # 将文本转为文件对象
                key = paramiko.RSAKey.from_private_key(cache)
                # 使用key登录
                ssh.connect(hostname=self.ip, port=self.port, username=self.username, pkey=key, timeout=self.timeout)
        except Exception as e:
            return {'status': 'ERROR', 'result': '连接目标主机超时！'}

        def exec_command(result_queue):
            stdin, stdout, stderr = ssh.exec_command(command)  # 没有超时机制，会一直等待命令响应
            result = stdout.read()
            error = stderr.read()
            if not error:
                # return {'status': 'SUCCESS', 'result': result}
                result_queue.put({'status': 'SUCCESS', 'result': result})
            else:
                result_queue.put({'status': 'ERROR', 'result': f'执行命令失败！错误：{e}'})
            ssh.close()

        result_queue = queue.Queue() # 将线程执行结果保存到队列中，以便主进程读取（下面）
        thread = threading.Thread(target=exec_command, args=(result_queue,)) # 创建一个线程对象
        thread.start()  # 启动线程
        thread.join(timeout=7)  # 阻塞主进程，等待子线程运行结束。如果线程在 7 秒内没有完成，主进程会继续执行，而不再等待子线程（只阻塞主进程7秒）。
        if thread.is_alive():  # 判断线程是否还在运行，如果返回True，说明线程未完成（超时）
            ssh.close()
            return {'status': 'ERROR', 'result': '执行命令超时！'}
        else:
            return result_queue.get()

    def scp(self, local_file, remote_file):
        def scp_task(result_queue):
            try:
                s = paramiko.Transport((self.ip, self.port)) # 无连接超时机制，使用线程解决
                if self.password:
                    s.connect(username=self.username, password=self.password)
                else:
                    cache = StringIO(self.key)  # 将文本转为文件对象
                    key = paramiko.RSAKey.from_private_key(cache)
                    s.connect(username=self.username, pkey=key, timeout=self.timeout)

                sftp = paramiko.SFTPClient.from_transport(s)
                sftp.put(local_file, remote_file)
                s.close()
                #return {'status': 'SUCCESS', 'result': '上传采集文件成功'}
                result_queue.put({'status': 'SUCCESS', 'result': '上传文件成功'})
            except Exception as e:
                #return {'status': 'ERROR', 'result': e}
                result_queue.put({'status': 'SUCCESS', 'result': '连接目标主机超时！'})

        result_queue = queue.Queue()
        thread = threading.Thread(target=scp_task, args=(result_queue,))
        thread.start()
        thread.join(timeout=7)
        if thread.is_alive():
            return {'status': 'ERROR', 'result': f'连接目标主机超时！'}
        else:
            return result_queue.get()


if __name__ == "__main__":
    ssh = SSH('192.168.1.72', 22, 'root', '123.com')
    result = ssh.scp('../files/host_collect.py', '/opt/host_collect.py')
    print(result)
    result = ssh.command('python /opt/host_collect.py')
    print(result)